from functools import cache, cached_property
from collections import UserString
from typing_extensions import Self


class Sequence(UserString):
    """A class for DNA/RNA/Protein sequence."""

    ATGCs = UserString.maketrans("ATGC", "TACG")
    NA2AA = {
        "TTT": "F",
        "TTC": "F",
        "TTA": "L",
        "TTG": "L",
        "TCT": "S",
        "TCC": "S",
        "TCA": "S",
        "TCG": "S",
        "TAT": "Y",
        "TAC": "Y",
        "TAA": "*",
        "TAG": "*",
        "TGT": "C",
        "TGC": "C",
        "TGA": "*",
        "TGG": "W",
        "CTT": "L",
        "CTC": "L",
        "CTA": "L",
        "CTG": "L",
        "CCT": "P",
        "CCC": "P",
        "CCA": "P",
        "CCG": "P",
        "CAT": "H",
        "CAC": "H",
        "CAA": "Q",
        "CAG": "Q",
        "CGT": "R",
        "CGC": "R",
        "CGA": "R",
        "CGG": "R",
        "ATT": "I",
        "ATC": "I",
        "ATA": "I",
        "ATG": "M",
        "ACT": "T",
        "ACC": "T",
        "ACA": "T",
        "ACG": "T",
        "AAT": "N",
        "AAC": "N",
        "AAA": "K",
        "AAG": "K",
        "AGT": "S",
        "AGC": "S",
        "AGA": "R",
        "AGG": "R",
        "GTT": "V",
        "GTC": "V",
        "GTA": "V",
        "GTG": "V",
        "GCT": "A",
        "GCC": "A",
        "GCA": "A",
        "GCG": "A",
        "GAT": "D",
        "GAC": "D",
        "GAA": "E",
        "GAG": "E",
        "GGT": "G",
        "GGC": "G",
        "GGA": "G",
        "GGG": "G",
    }

    @cached_property
    def size(self) -> int:
        """Return the size of the sequence."""
        return len(self)

    @cache
    def reverse(self) -> Self:
        """Return the reverse of the sequence."""
        return self[::-1]

    @cache
    def reverse_complement(self) -> Self:
        """Return the complement of the sequence."""
        return self.reverse().translate(self.ATGCs)

    @cache
    def protein(self) -> Self:
        """Return the protein sequence."""
        bases = []
        for i in range(0, self.size, 3):
            bases.append(self.NA2AA.get(self[i : i + 3], "X"))
        return self.__class__("".join(bases))

    def substitute(self, pos: int, ref: Self, alt: Self):
        """Return the sequence with the position substituted. pos: 1-based."""
        if pos:
            return self[: pos - 1] + alt + self[pos + ref.size - 1 :]
        return alt[ref.size :] + self

    def delete(self, start: int, end: int) -> Self:
        """Return the sequence with the position deleted. start,end: 1-based."""
        return self[: start - 1] + self[end:]

    def dup_unit(self) -> Self:
        """Return the sequence with the unit duplicated."""
        seq_len = self.size
        for unit_len in filter(lambda x: seq_len % x == 0, range(1, seq_len // 2 + 1)):
            unit = self[:unit_len]
            num = seq_len // unit_len
            if unit * num == self:
                return unit
        return self

    def difference_simple(self, other: Self) -> int:
        """Return the sequence with the first difference position. 1-based."""
        for i, (x, y) in enumerate(zip(self, other)):
            if x != y:
                return i + 1
        return min(self.size, other.size) + 1

    def difference(self, other: Self) -> tuple[int, int, int]:
        """Return the sequence with the first difference position and the each last difference position. 1-based."""
        left_len = self.difference_simple(other) - 1
        right_len = self.reverse().difference_simple(other.reverse()) - 1
        min_len = min(self.size, other.size)
        if left_len + right_len > min_len:
            right_len = min_len - left_len
        return left_len + 1, self.size - right_len, other.size - right_len


class MTSequence(Sequence):
    NA2AA = {
        "TTT": "F",
        "TTC": "F",
        "TTA": "L",
        "TTG": "L",
        "TCT": "S",
        "TCC": "S",
        "TCA": "S",
        "TCG": "S",
        "TAT": "Y",
        "TAC": "Y",
        "TAA": "*",
        "TAG": "*",
        "TGT": "C",
        "TGC": "C",
        "TGA": "W",
        "TGG": "W",
        "CTT": "L",
        "CTC": "L",
        "CTA": "L",
        "CTG": "L",
        "CCT": "P",
        "CCC": "P",
        "CCA": "P",
        "CCG": "P",
        "CAT": "H",
        "CAC": "H",
        "CAA": "Q",
        "CAG": "Q",
        "CGT": "R",
        "CGC": "R",
        "CGA": "R",
        "CGG": "R",
        "ATT": "I",
        "ATC": "I",
        "ATA": "M",
        "ATG": "M",
        "ACT": "T",
        "ACC": "T",
        "ACA": "T",
        "ACG": "T",
        "AAT": "N",
        "AAC": "N",
        "AAA": "K",
        "AAG": "K",
        "AGT": "S",
        "AGC": "S",
        "AGA": "*",
        "AGG": "*",
        "GTT": "V",
        "GTC": "V",
        "GTA": "V",
        "GTG": "V",
        "GCT": "A",
        "GCC": "A",
        "GCA": "A",
        "GCG": "A",
        "GAT": "D",
        "GAC": "D",
        "GAA": "E",
        "GAG": "E",
        "GGT": "G",
        "GGC": "G",
        "GGA": "G",
        "GGG": "G",
    }
